package middleware

import (
	"context"
	"sync"

	"golang.org/x/time/rate"
)

// RateLimiter implements rate limiting
type RateLimiter struct {
	limiters map[string]*rate.Limiter
	mutex    sync.RWMutex
}

// NewRateLimiter creates a new rate limiter
func NewRateLimiter() *RateLimiter {
	return &RateLimiter{
		limiters: make(map[string]*rate.Limiter),
	}
}

// SetLimit sets rate limit for a provider
func (r *RateLimiter) SetLimit(provider string, rps float64) {
	r.mutex.Lock()
	defer r.mutex.Unlock()

	r.limiters[provider] = rate.NewLimiter(rate.Limit(rps), int(rps))
}

// Wait waits for rate limit
func (r *RateLimiter) Wait(ctx context.Context, provider string) error {
	r.mutex.RLock()
	limiter, ok := r.limiters[provider]
	r.mutex.RUnlock()

	if !ok {
		return nil // No limit set for this provider
	}

	return limiter.Wait(ctx)
}
